package com.yuncheng.mcp.tool;

import jakarta.annotation.PostConstruct;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.*;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.metadata.ToolMetadata;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;


@Component
@Primary
public class DynamicToolCallbackProvider implements ToolCallbackProvider {

    private final Map<String, DynamicTool> toolRegistry = new ConcurrentHashMap<>();

    // 工具配置缓存，替代数据库方案
    private final Map<String, ToolConfig> toolConfigCache = new ConcurrentHashMap<>();

    /**
     * 初始化方法，在应用启动时加载工具
     */
    @PostConstruct
    public void init() {
        // 从缓存加载预定义工具
        loadPredefinedTools();
        System.out.println(">>> 已加载 " + toolRegistry.size() + " 个动态工具");
    }

    /**
     * 加载预定义工具示例
     */
    private void loadPredefinedTools() {
        // 这里可以添加一些默认工具
        addToolConfig(new ToolConfig(
                "power",
                "计算一个数的幂",
                "Math.pow(#base, #exponent)",
                Map.of("base", "java.lang.Double", "exponent", "java.lang.Double")
        ));

        addToolConfig(new ToolConfig(
                "sqrt",
                "计算平方根",
                "Math.sqrt(#number)",
                Map.of("number", "java.lang.Double")
        ));

        // 加载所有缓存的工具配置
        toolConfigCache.values().forEach(this::registerToolFromConfig);
    }

    /**
     * 添加工具配置到缓存
     */
    public void addToolConfig(ToolConfig config) {
        toolConfigCache.put(config.getName(), config);
        registerToolFromConfig(config);
    }

    /**
     * 从配置创建并注册工具
     */
    private void registerToolFromConfig(ToolConfig config) {
        try {
            Map<String, Class<?>> paramTypes = parseParameterTypes(config.getParameterTypes());

            DynamicTool tool = new SimpleDynamicTool(
                    config.getName(),
                    config.getDescription(),
                    config.getExpression(),
                    paramTypes
            );

            registerTool(tool);
        } catch (Exception e) {
            System.err.println("无法注册工具 " + config.getName() + ": " + e.getMessage());
        }
    }

    /**
     * 解析参数类型
     */
    private Map<String, Class<?>> parseParameterTypes(Map<String, String> typeStrings) {
        Map<String, Class<?>> result = new ConcurrentHashMap<>();
        typeStrings.forEach((name, typeStr) -> {
            try {
                Class<?> type = Class.forName(typeStr);
                result.put(name, type);
            } catch (ClassNotFoundException e) {
                throw new RuntimeException("无法加载参数类型: " + typeStr, e);
            }
        });
        return result;
    }

    /**
     * 获取所有工具配置
     */
    public List<ToolConfig> getAllToolConfigs() {
        return new ArrayList<>(toolConfigCache.values());
    }

    /**
     * 根据名称获取工具配置
     */
    public ToolConfig getToolConfig(String name) {
        return toolConfigCache.get(name);
    }

    /**
     * 删除工具配置
     */
    public void removeToolConfig(String name) {
        toolConfigCache.remove(name);
        unregisterTool(name);
    }

    /**
     * 更新工具配置
     */
    public void updateToolConfig(ToolConfig config) {
        removeToolConfig(config.getName());
        addToolConfig(config);
    }

    /**
     * 注册一个动态工具
     * @param tool 要注册的工具
     */
    public void registerTool(DynamicTool tool) {
        toolRegistry.put(tool.getName(), tool);
    }

    /**
     * 注销一个动态工具
     * @param toolName 要注销的工具名称
     */
    public void unregisterTool(String toolName) {
        toolRegistry.remove(toolName);
    }

    /**
     * 获取所有注册的工具
     * @return 工具列表
     */
    public List<DynamicTool> getAllTools() {
        return new ArrayList<>(toolRegistry.values());
    }

    /**
     * 获取工具回调数组
     */
    @Override
    public ToolCallback[] getToolCallbacks() {
        List<ToolCallback> toolCallbacks = new ArrayList<>();

        // 将每个动态工具转换为ToolCallback
        for (DynamicTool tool : toolRegistry.values()) {
            toolCallbacks.add(createToolCallback(tool));
        }

        return toolCallbacks.toArray(new ToolCallback[0]);
    }

    /**
     * 为动态工具创建ToolCallback
     */
    private ToolCallback createToolCallback(DynamicTool tool) {
        return new ToolCallback() {
            @Override
            public ToolDefinition getToolDefinition() {
                // 创建工具定义
                return ToolDefinition.builder()
                        .name(tool.getName())
                        .description(tool.getDescription())
                        .inputSchema(tool.generateInputSchema())
                        .build();
            }

            @Override
            public String call(String toolInput) {
                System.out.println("call toolInput: " + toolInput);
                try {
                    // 解析输入JSON为参数Map
                    Map<String, Object> arguments = tool.parseInput(toolInput);
                    Object result = tool.execute(arguments);
                    return tool.formatOutput(result);
                } catch (Exception e) {
                    return "{\"error\": \"" + e.getMessage() + "\"}";
                }
            }

            // 其他注释掉的方法保持不变
        };
    }

    /**
     * 工具配置类
     */
    public static class ToolConfig {
        private String name;
        private String description;
        private String expression;
        private Map<String, String> parameterTypes; // 参数名 -> 类型全限定名

        public ToolConfig() {}

        public ToolConfig(String name, String description, String expression, Map<String, String> parameterTypes) {
            this.name = name;
            this.description = description;
            this.expression = expression;
            this.parameterTypes = parameterTypes;
        }

        // Getters and setters
        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        public String getDescription() {
            return description;
        }

        public void setDescription(String description) {
            this.description = description;
        }

        public String getExpression() {
            return expression;
        }

        public void setExpression(String expression) {
            this.expression = expression;
        }

        public Map<String, String> getParameterTypes() {
            return parameterTypes;
        }

        public void setParameterTypes(Map<String, String> parameterTypes) {
            this.parameterTypes = parameterTypes;
        }
    }
}
